package demo.flyweight;


import java.nio.ByteBuffer;
import java.util.ArrayDeque;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

/**
 * @author : feixiang.li
 * @since : 2025-09-17 20:48
 */
public class BufferPool {

    private final int totalSize;

    private final int slotSize;

    private int free;

    private final ArrayDeque<ByteBuffer> slotQueue = new ArrayDeque<>();
    private final ArrayDeque<Condition> waitQueue = new ArrayDeque<>();

    private final Lock lock = new ReentrantLock();

    public BufferPool(int totalSize, int slotSize) {
        this.totalSize = totalSize;
        this.slotSize = slotSize;
        this.free = totalSize;
    }


    public ByteBuffer allocate(int size, int timeout) throws InterruptedException {
        if (size > totalSize || size < 0) {
            throw new IllegalArgumentException("申请容量为空");
        }
        lock.lock();
        try {
            // 如果请求的容量和slotSize一致，并且槽的队列中有数据，则直接从队列中获取·
            ByteBuffer buffer = check(size);
            if (buffer != null) {
                return buffer;
            }
            Condition condition = lock.newCondition();
            waitQueue.addLast(condition);
            long remain = timeout;
            try {
                while (true) {
                    long start = System.currentTimeMillis();
                    boolean weekUp = condition.await(remain, TimeUnit.MILLISECONDS);
                    if (!weekUp) {
                        throw new InterruptedException();
                    }
                    buffer = check(size);
                    if (buffer != null) {
                        return buffer;
                    }
                    remain = remain - (System.currentTimeMillis() - start);
                }
            } finally {
                waitQueue.remove(condition);
            }
        } finally {
            // 自己申请完成后，开始 唤醒等待的线程
            if (!waitQueue.isEmpty() && !(free == 0 && slotQueue.isEmpty())) {
                // 这里只是唤醒。等待线程自己会把自己移除掉
                waitQueue.peekFirst().signal();
            }
            lock.unlock();
        }
    }


    public ByteBuffer check(int size) {
        if (size == slotSize && !slotQueue.isEmpty()) {
            return slotQueue.pollFirst();
        }
        // 如果请求的容量。> 剩余的容量+槽队列中数据，则返回nul
        if (free + slotQueue.size() * slotSize >= size) {
            freeUp(size);
            free -= size;
            return ByteBuffer.allocate(size);
        }
        return null;
    }

    private void freeUp(int size) {
        while (free < size && !slotQueue.isEmpty()) {
            free += slotQueue.pollFirst().capacity();
        }
    }


    public void dealocate(ByteBuffer buffer) {
        lock.lock();
        try {
            if (buffer.capacity() == slotSize) {
                buffer.clear();
                slotQueue.addLast(buffer);
            } else {
                free += buffer.capacity();
            }
            // 唤醒等待的线程
            if (!waitQueue.isEmpty()) {
                waitQueue.peekFirst().signal();
            }
        } finally {
            lock.unlock();
        }
    }
}
